// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <thrift/transport/TVirtualTransport.h>
#include <pollux/dwio/common/buffered_input.h>

namespace kumo::pollux::parquet::thrift {

class ThriftTransport
    : public apache::thrift::transport::TVirtualTransport<ThriftTransport> {
 public:
  virtual uint32_t read(uint8_t* outputBuf, uint32_t len) = 0;
  virtual ~ThriftTransport() = default;
};

class ThriftStreamingTransport : public ThriftTransport {
 public:
  ThriftStreamingTransport(
      dwio::common::SeekableInputStream* inputStream,
      const char*& bufferStart,
      const char*& bufferEnd)
      : inputStream_(inputStream),
        bufferStart_(bufferStart),
        bufferEnd_(bufferEnd) {
    POLLUX_CHECK_NOT_NULL(inputStream_);
    POLLUX_CHECK_NOT_NULL(bufferStart_);
    POLLUX_CHECK_NOT_NULL(bufferEnd_);
  }

  uint32_t read(uint8_t* outputBuf, uint32_t len) {
    uint32_t bytesToRead = len;
    while (bytesToRead > 0) {
      if (bufferEnd_ == bufferStart_) {
        int32_t size;
        if (!inputStream_->Next(
                reinterpret_cast<const void**>(&bufferStart_), &size)) {
          POLLUX_FAIL("Reading past the end of the stream");
        }
        bufferEnd_ = bufferStart_ + size;
      }

      uint32_t bytesToReadInBuffer =
          std::min<uint32_t>(bufferEnd_ - bufferStart_, bytesToRead);
      memcpy(outputBuf, bufferStart_, bytesToReadInBuffer);
      bufferStart_ += bytesToReadInBuffer;
      bytesToRead -= bytesToReadInBuffer;
      outputBuf += bytesToReadInBuffer;
    }

    return len;
  }

 private:
  dwio::common::SeekableInputStream* inputStream_;
  const char*& bufferStart_;
  const char*& bufferEnd_;
};

class ThriftBufferedTransport : public ThriftTransport {
 public:
  ThriftBufferedTransport(const void* inputBuf, uint64_t len)
      : ThriftTransport(),
        inputBuf_(reinterpret_cast<const uint8_t*>(inputBuf)),
        size_(len),
        offset_(0) {}

  uint32_t read(uint8_t* outputBuf, uint32_t len) {
    DWIO_ENSURE(offset_ + len <= size_);
    memcpy(outputBuf, inputBuf_ + offset_, len);
    offset_ += len;
    return len;
  }

 private:
  const uint8_t* inputBuf_;
  const uint64_t size_;
  uint64_t offset_;
};

} // namespace kumo::pollux::parquet::thrift
